/*===============================================================================================
 Copyright©  Junjun Yan. All rights reserved.
 File     : 向量求和.cpp
 Author   : shuaicaijunjun@gmail.com
 Created  : 2020/11/27
 Summary  : 1.Read_Vector函数实现读入向量并分发给各个进程的功能，其中调用的MPI_Scatter函数要求使用
			  块划分法，并且向量维数n必须被进程数comm_z整除
			2.Print_Vector函数实现将各进程的向量聚合卡里并输出的功能，其中调用的MPI_Gather的参数与
			  MPI_Scatter的参数相似，Print_Vector的参数也与Read_Vector函数相似
 ===============================================================================================*/
 
#include<iostream>
#include<string>
#include"mpi.h"
 
using namespace std;
 
void Read_Vector(
	int n,				/*输入向量的维数，即总维数*/
	int local_n,		/*输出向量的维数，即每个进程的维数*/
	double vector[],	/*输入向量的地址*/
	string vector_name, /*输入向量的名称*/
	int my_rank,		/*进程号*/
	MPI_Comm comm		/*通信子*/) {
 
	double* a = NULL;
	
	// 主进程开辟空间读入向量，并调用scatter函数将向量分发至各个进程
	if (my_rank == 0) {
		a = new double[n];
		cout << "Enter the vector " << vector_name << ":" << endl;
		for (int i = 0; i < n; i++) {
			cin >> a[i];
		}
		//使用scatter函数分发时只能使用块划分，并且必须满足n能整除comm_sz
		MPI_Scatter(a, local_n,	MPI_DOUBLE, vector, local_n, MPI_DOUBLE, 0, comm);
		delete[] a;
	}
	//其余进程接收向量
	else {
		MPI_Scatter(a, local_n, MPI_DOUBLE, vector, local_n, MPI_DOUBLE, 0, comm);
	}
}
 
void vector_sum(double local_a[], double local_b[], int local_n, double local_c[]) {
	/*向量求和，和串行的没有什么区别*/
	for (int i = 0; i < local_n; i++) {
		local_c[i] = local_a[i] + local_b[i];
	}
}
 
void Print_vector(int n, int local_n, double local_c[], string title, int my_rank, MPI_Comm comm) {
	double* c = NULL;
	// 主进程开辟空间，将其余进程计算好的数据接收并聚集为一个向量
	if (my_rank == 0) {
		c = new double[n];
		MPI_Gather(local_c, local_n, MPI_DOUBLE, c, local_n, MPI_DOUBLE, 0, comm);
		cout << title << endl;
		for (int i = 0; i < n; i++) {
			cout << c[i] << " ";
		}
		cout << endl;
        delete[] c;
	}
	// 其余进程将计算结果发送至0号进程
	else {
		MPI_Gather(local_c, local_n, MPI_DOUBLE, c, local_n, MPI_DOUBLE, 0, comm);
	}
}
 
int main() {
	int comm_sz;	// 进程数
	int my_rank;	// 进程号
	double* local_a, * local_b, * local_c;	// 每个进程的向量
	int n,local_n;	//向量维数以及每个进程的向量的维数
 
	MPI_Init(NULL, NULL);
 
	MPI_Comm_size(MPI_COMM_WORLD, &comm_sz);
	MPI_Comm_rank(MPI_COMM_WORLD, &my_rank);
 
	// 0号进程从控制台接收需要创建的向量维数 
	if (my_rank == 0) {
		cout << "Enter n:" << endl;
		cin >> n;
	}
	// 将接收到的n广播至各个进程
	MPI_Bcast(&n, 1, MPI_INT, 0, MPI_COMM_WORLD);
 
	// 计算每个进程所需要的向量维数
	local_n = n / comm_sz;
 
	// 按照接收到的n值来为每个进程动态创建数组大小
	local_a = new double[local_n];
	local_b = new double[local_n];
	local_c = new double[local_n];
 
	// 获取向量a和向量b的数据
	Read_Vector(n, local_n, local_a, "a", my_rank, MPI_COMM_WORLD);
	Read_Vector(n, local_n, local_b, "b", my_rank, MPI_COMM_WORLD);
 
	//每个进程执行向量相加
	vector_sum(local_a, local_b, local_n, local_c);
 
	// 输出计算结果
	Print_vector(n, local_n, local_c, "The vector c:", my_rank, MPI_COMM_WORLD);
 
	MPI_Finalize();
 
	return 0;
}